/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024.
 * Description: reclaim memory requested by squashfs
 * Author: Pang LiYuan
 * Create: 2024-08-19
 */

#include <linux/proc_fs.h>
#include <linux/fs.h>
#include <linux/seq_file.h>
#include <linux/slab.h>
#include <linux/shrinker.h>
#include <linux/module.h>

#include "squashfs_fs_sb.h"
#include "squashfs.h"
#include "page_actor.h"
#include "rtos_shrink.h"

enum {
	SHRINK_DISABLED = 0, /* don't shrink memory anyway */
	SHRINK_DIRECT = 1, /* shrink memory right now */
	SHRINK_MEM_RECLAIM = 2, /* shrink memory in kernel memory reclaim process */
};

/* to find all mounted suqashfs */
LIST_HEAD(squashfs_infos);
DEFINE_MUTEX(squashfs_infos_lock);
atomic_long_t squashfs_shrink_cnt;

int less_frag;
static int shrink_mode;

#ifndef MODULE
static int __init get_less_frag(char *str)
{
	get_option(&str, &less_frag);
	return 0;
}
early_param("less_frag", get_less_frag);
#else
module_param(less_frag, int, 0400);
MODULE_PARM_DESC(less_frag, "enable fragment cache size optimize");
#endif  /* MODULE */

int squashfs_cache_shrink_entry(struct squashfs_cache *cache)
{
	int i, j, reclaims = 0;
	struct squashfs_cache_entry *entry;

	spin_lock(&cache->lock);

	for (i = 0; i < cache->entries && cache->unused; i++) {
		entry = &cache->entry[i];
		/*
		 * refcount == 1 means entry in use
		 * valid == 0 means entry has been taken back
		 */
		if (entry->refcount || !entry->valid)
			continue;

		if (entry->data) {
			for (j = 0; j < cache->pages; j++)
				kfree(entry->data[j]);
			kfree(entry->data);
			kfree(entry->actor);
			atomic_long_dec(&squashfs_shrink_cnt);
		}
		entry->data = NULL;
		entry->actor = NULL;
		entry->valid = 0;
		entry->block = SQUASHFS_INVALID_BLK;
		cache->unused--;
		cache->valid_cnt--;
		reclaims++;
	}

	cache->miss = 0;
	spin_unlock(&cache->lock);
	return reclaims;
}

int squashfs_cache_refill_entry(struct squashfs_cache *cache, struct squashfs_cache_entry *entry)
{
	int j;

	entry->data = kcalloc(cache->pages, sizeof(void *), GFP_KERNEL);
	if (entry->data == NULL) {
		ERROR("Failed to allocate %s cache entry\n", cache->name);
		return -ENOMEM;
	}

	for (j = 0; j < cache->pages; j++) {
		entry->data[j] = kmalloc(PAGE_SIZE, GFP_KERNEL);
		if (entry->data[j] == NULL) {
			ERROR("Failed to allocate %s buffer\n", cache->name);
			goto clean_page;
		}
	}

	entry->actor = squashfs_page_actor_init(entry->data,
					cache->pages, 0);
	if (entry->actor == NULL) {
		ERROR("Failed to allocate %s cache entry\n", cache->name);
		goto clean_page;
	}

	atomic_long_inc(&squashfs_shrink_cnt);
	return 0;

clean_page:
	for (j = 0; j < cache->pages && entry->data[j]; j++)
		kfree(entry->data[j]);
	kfree(entry->data);
	return -ENOMEM;
}

void squashfs_cache_increase_entry(struct squashfs_cache *cache)
{
	int i, n;

	i = cache->next_blk;
	for (n = 0; n < cache->entries; n++) {
		if (cache->entry[i].refcount == 0 && !cache->entry[i].valid) {
			cache->entry[i].valid = 1;
			cache->unused++;
			cache->valid_cnt++;
			break;
		}

		i = (i + 1) % cache->entries;
	}
}

static unsigned long squashfs_shrink_mem(void)
{
	unsigned long tot_cnt = 0;
	struct squashfs_sb_info *msblk;

	if (atomic_long_read(&squashfs_shrink_cnt) == 0)
		return SHRINK_STOP;

	mutex_lock(&squashfs_infos_lock);
	list_for_each_entry(msblk, &squashfs_infos, infos_list) {
		if (msblk->thread_ops->shrink)
			tot_cnt += msblk->thread_ops->shrink(msblk);

		tot_cnt += squashfs_cache_shrink_entry(msblk->read_page);

		/* if no fragments, this should be NULL */
		if (msblk->fragment_cache)
			tot_cnt += squashfs_cache_shrink_entry(msblk->fragment_cache);
	}
	mutex_unlock(&squashfs_infos_lock);
	return tot_cnt;
}

static ssize_t squashfs_write_proc(struct file *file, const char __user *buffer,
									size_t size, loff_t *ppos)
{
	char tmp;

	if (size > 2)
		return -EINVAL;

	if (copy_from_user(&tmp, buffer, 1))
		return -EFAULT;

	/*
	 * four valid options from 0-3
	 * 0:disabled
	 * 1:only start shrink immediately
	 * 2:only enable memory reclamation process
	 * 3:both 1 & 2
	 */
	if (tmp < '0' || tmp > '3')
		return -EINVAL;

	shrink_mode = tmp - '0';
	if (shrink_mode & SHRINK_DIRECT)
		squashfs_shrink_mem();

	return size;
}

static int squashfs_proc_show(struct seq_file *m, void *v)
{
	seq_printf(m, "current shrink mode:%d\n", shrink_mode);
	return 0;
}

static int squashfs_proc_open(struct inode *inode, struct file *file)
{
	if (!capable(CAP_SYS_ADMIN))
		return -EACCES;

	return single_open(file, squashfs_proc_show, NULL);
}

static struct proc_ops squashfs_shrink_fops = {
	.proc_open = squashfs_proc_open,
	.proc_release = single_release,
	.proc_read = seq_read,
	.proc_lseek = seq_lseek,
	.proc_write = squashfs_write_proc,
};

static unsigned long squashfs_shrink_scan(struct shrinker *shrink,
				struct shrink_control *sc)
{
	if (shrink_mode & SHRINK_MEM_RECLAIM)
		return squashfs_shrink_mem();

	return 0;
}

static unsigned long squashfs_shrink_count(struct shrinker *shrink,
				struct shrink_control *sc)
{
	if (shrink_mode & SHRINK_MEM_RECLAIM)
		return atomic_long_read(&squashfs_shrink_cnt);

	return 0;
}

static struct shrinker squashfs_shrinker_info = {
	.scan_objects = squashfs_shrink_scan,
	.count_objects = squashfs_shrink_count,
	.seeks = DEFAULT_SEEKS,
	.flags = SHRINKER_RTOS_HIGHPRIO,
};

int squashfs_shrink_init(void)
{
	int err = 0;

	if (proc_create("squashfs_shrink_mem", 0666, NULL, &squashfs_shrink_fops) == NULL) {
		pr_err("squashfs shrink mem proc create failed\n");
		return -ENOMEM;
	}

	err = register_shrinker(&squashfs_shrinker_info);
	if (err) {
		remove_proc_entry("squashfs_shrink_mem", NULL);
		return err;
	}

	return 0;
}

void squashfs_shrink_exit(void)
{
	unregister_shrinker(&squashfs_shrinker_info);
	remove_proc_entry("squashfs_shrink_mem", NULL);
}
